import os
from os.path import join, isdir, isfile
import numpy as np
import torch
from transforms3d.euler import euler2mat
from transforms3d.axangles import mat2axangle
from torch_utils import batch_rodrigues
import open3d as o3d


def get_bbox(pts):
    x_min, y_min, z_min = list(pts.min(axis=0))
    x_max, y_max, z_max = list(pts.max(axis=0))
    return [x_min, y_min, z_min, x_max, y_max, z_max]


def align(root_dir, folder_name, obj2world_R):

    chair_list = []

    data_dir = join(root_dir, folder_name)
    save_dir = join(root_dir, folder_name.replace("_notalign", ""))
    for scene_name in os.listdir(data_dir):
        scene_dir = join(data_dir, scene_name)
        if not isdir(scene_dir):
            continue

        scene_mesh_path = join(scene_dir, "scene_mesh.obj")
        motion_dir = join(scene_dir, "h1_kinematic_motions")
        if (not isfile(scene_mesh_path)) or (not isdir(motion_dir)):
            continue
        scene_mesh = o3d.io.read_triangle_mesh(scene_mesh_path)
        pts = np.float32(scene_mesh.vertices)  # (N, 3)
        faces = np.int32(scene_mesh.triangles)
        bbox = get_bbox(pts)
        if bbox[3] - bbox[0] < 1.5:
            chair_list.append(scene_name)
        cx, cy = (bbox[0] + bbox[3]) / 2, (bbox[1] + bbox[4]) / 2
        obj2world = np.eye(4)
        obj2world[:3, :3] = obj2world_R
        obj2world[:3, 3] = np.float32([cx, cy, 0])

        motion_save_dir = join(save_dir, scene_name, "h1_kinematic_motions")
        os.makedirs(motion_save_dir, exist_ok=True)
        # save aligned scene mesh
        pts = (pts - obj2world[:3, 3].reshape(1, 3)) @ obj2world[:3, :3]
        aligned_scene_mesh = o3d.geometry.TriangleMesh(vertices=o3d.utility.Vector3dVector(pts), triangles=o3d.utility.Vector3iVector(faces))
        o3d.io.write_triangle_mesh(join(save_dir, scene_name, "scene_mesh.obj"), aligned_scene_mesh)
        # save aligned h1 motion
        for h1_motion_name in os.listdir(motion_dir):
            motion = np.load(join(motion_dir, h1_motion_name), allow_pickle=True)["arr_0"].item()
            aligned_motion = {
                "joint_angles": motion["joint_angles"].copy(),
            }
            N = motion["joint_angles"].shape[0]
            T = np.eye(4).reshape(1, 4, 4).repeat(N, 0)
            T[:, :3, 3] = motion["global_translations"]
            T[:, :3, :3] = batch_rodrigues(torch.from_numpy(motion["global_rotations"])).detach().cpu().numpy()
            T = np.linalg.inv(obj2world) @ T
            aligned_motion["global_translations"] = T[:, :3, 3]
            rotvecs = []
            for i in range(N):
                axis, angle = mat2axangle(T[i, :3, :3], unit_thresh=1e-3)
                rotvecs.append(axis * angle)
            aligned_motion["global_rotations"] = np.float32(rotvecs)
            aligned_motion["link_global_poses"] = {}
            for link_name in motion["link_global_poses"]:
                aligned_motion["link_global_poses"][link_name] = np.linalg.inv(obj2world) @ motion["link_global_poses"][link_name]
            np.savez(join(motion_save_dir, h1_motion_name), aligned_motion)

    print("chairs:")
    print(len(chair_list))
    print(chair_list)


if __name__ == "__main__":

    ##################### change here ########################################
    root_dir = "./"
    folder_name = "copy_rotation_UniHSI_retargeted_data_augmented_chair_sit_notalign"
    ##########################################################################

    assert folder_name.endswith("_notalign")

    obj2world_R = np.eye(3)
    if folder_name.find("ROAM") > -1:
        obj2world_R = euler2mat(0, 0, 30/180*np.pi)

    align(root_dir, folder_name, obj2world_R)
